﻿//
// Created by QU on 24-9-28.
//

#ifndef STRESSCLIENT_H
#define STRESSCLIENT_H

#include <boost/asio.hpp>
#include <chrono>
#include <atomic>
#include <random>
#include <string>
#include <vector>
#include <future>

#include "node/Node.h"
#include "server/logic/MSGID.h"

class StressClient
{
public:
    struct Config
    {
        int thread_count = 10;
        int messages_per_thread = 10000;
        int small_message_min = 10;
        int small_message_max = 1024;
        int large_message_min = 1024;
        int large_message_max = Node::MAX_PAYLOAD_SIZE;
        int multiple_message_count = 10;
        int multiple_message_min = 5;
        int multiple_message_max = 8192;
        int split_message_min = 5;
        int split_message_max = 8192;
        int delay_between_messages_min = 0;
        int delay_between_messages_max = 5;
        std::chrono::seconds total_test_duration{ 300 };
    };

    StressClient(const std::string& host, const std::string& port, const Config& config = Config {});

    void startLoad();

private:
    void runClientThread(std::atomic<bool>& should_stop, std::chrono::steady_clock::time_point start_time);
    void connect(boost::asio::ip::tcp::socket& socket) const;
    void sendRandomMessage(boost::asio::ip::tcp::socket& socket, int message_type);
    void sendSmallMessage(boost::asio::ip::tcp::socket& socket);
    void sendLargeMessage(boost::asio::ip::tcp::socket& socket);
    void sendMultipleMessages(boost::asio::ip::tcp::socket& socket);
    void sendSplitMessage(boost::asio::ip::tcp::socket& socket);
    void transmitMessage(boost::asio::ip::tcp::socket& socket, size_t payload_size, const std::string& message_type);
    std::vector<std::byte> constructMessage(size_t payload_size);
    void delayBetweenMessages(std::default_random_engine& gen);
    size_t randomSize(int min, int max);
    void receiveEcho(boost::asio::ip::tcp::socket& socket);

    std::string m_host;
    std::string m_port;
    Config m_config;
    std::mt19937 m_gen { std::random_device { }() };
    std::atomic<size_t> m_total_sent_messages{0};
    std::atomic<size_t> m_total_received_messages{0};
};

#endif // STRESSCLIENT_H